# -*- coding: UTF-8 -*-
# 为方便测试，请统一使用 numpy、pandas、sklearn 三种包，如果实在有特殊需求，请单独跟助教沟通
import numpy as np
import pandas as pd
from sklearn import pipeline
from sklearn.impute import SimpleImputer
from sklearn.metrics import f1_score
from sklearn import svm
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier as DTC, export_graphviz
from sklearn.model_selection import GridSearchCV
from sklearn.decomposition import PCA
import argparse

np.set_printoptions(suppress=True)
# 设定随机数种子，保证代码结果可复现
np.random.seed(1024)


class Model:
    """
    要求：
        1. 需要有__init__、train、predict三个方法，且方法的参数应与此样例相同
        2. 需要有self.X_train、self.y_train、self.X_test三个实例变量，请注意大小写
        3. 如果划分出验证集，请将实例变量命名为self.X_valid、self.y_valid
    """

    # 模型初始化，数据预处理，仅为示例
    def __init__(self, train_path, test_path):
        df = pd.read_csv(train_path, encoding='gbk', index_col='id')
        df_test = pd.read_csv(test_path, encoding='gbk', index_col='id')
        # print(df.shape[1])
        # RBP4列缺失值实在是太多了，删除
        df = df.drop(labels=['RBP4'], axis=1)
        df_test = df_test.drop(labels=['RBP4'], axis=1)

        # 处理缺失值
        df = df.dropna(how='all')
        df_test = df_test.dropna(how='all')
        # 离散值用众数代替缺失值
        for i in range(23):
            df.iloc[:, i].fillna(df.iloc[:, i].mode()[0], inplace=True)
            df_test.iloc[:, i].fillna(df_test.iloc[:, i].mode()[0], inplace=True)
        # print(df_train.iloc[:,81].name)
        for i in range(48, 82):
            df.iloc[:, i].fillna(df.iloc[:, i].mode()[0], inplace=True)
            df_test.iloc[:, i].fillna(df_test.iloc[:, i].mode()[0], inplace=True)
        df.loc[:, '孕次'].fillna(df.loc[:, '孕次'].mode()[0], inplace=True)
        df.loc[:, '产次'].fillna(df.loc[:, '产次'].mode()[0], inplace=True)
        df_test.loc[:, '孕次'].fillna(df_test.loc[:, '孕次'].mode()[0], inplace=True)
        df_test.loc[:, '产次'].fillna(df_test.loc[:, '产次'].mode()[0], inplace=True)

        cols = []
        # 将未处理的列的缺失值用平均值填充
        data_preprocessing = SimpleImputer(strategy="mean")
        df_train, df_valid, self.y_train, self.y_valid = train_test_split(df, df['label'], test_size=0.2)
        del(df_train['label'])
        del(df_valid['label'])
        self.X_train = data_preprocessing.fit_transform(df_train)
        self.X_valid = data_preprocessing.fit_transform(df_valid)
        # print('df_test:{}'.format(df_test.shape[1]))
        self.X_test = data_preprocessing.fit_transform(df_test)
        # print('df_test:{}'.format(self.X_test.shape[1]))
        # 对部分列归一化
        # print(self.X_train[:,47])
        l_age = self.X_train[:, 23]
        self.X_train[:, 23] = (l_age - np.mean(l_age)) / np.std(l_age)
        l_age_2 = self.X_valid[:, 23]
        self.X_valid[:, 23] = (l_age_2 - np.mean(l_age_2)) / np.std(l_age_2)
        l_age_3 = self.X_test[:, 23]
        self.X_test[:, 23] = (l_age_3 - np.mean(l_age_3)) / np.std(l_age_3)

        for i in range(26, 48):
            l_tmp = self.X_train[:, i]
            l_tmp_2 = self.X_valid[:,i]
            l_tmp_3 = self.X_test[:, i]
            self.X_train[:, i] = (l_tmp - np.mean(l_tmp)) / np.std(l_tmp)
            self.X_valid[:, i] = (l_tmp_2 - np.mean(l_tmp_2)) / np.std(l_tmp_2)
            self.X_test[:, i] = (l_tmp_3 - np.mean(l_tmp_3)) / np.std(l_tmp_3)
        # print(self.X_test.shape[1])
        # # np.savetxt('./test.csv', self.X_train, delimiter=',',fmt='%.04f')
        # print(self.X_train)
        # https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html
        self.classification_model = svm.SVC(C = 2, kernel='rbf', gamma="auto")
        self.df_predict = pd.DataFrame(index=df_test.index)

    # 模型训练，输出训练集f1_score
    def train(self):
        self.classification_model.fit(self.X_train, self.y_train)
        y_valid_pred = self.classification_model.predict(self.X_valid)
        return f1_score(self.y_valid, y_valid_pred)

    # 模型测试，输出测试集预测结果，要求此结果为DataFrame格式，可以通过to_csv方法保存为Kaggle的提交文件
    def predict(self):
        y_test_pred = self.classification_model.predict(self.X_test)
        self.df_predict['Predicted'] = y_test_pred
        return self.df_predict


# 以下部分请勿改动！
if __name__ == '__main__':
    # 解析输入参数。在终端执行以下语句即可运行此代码： python f_model.py --train_path "f_train.csv" --test_path "f_test.csv"
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_path", type=str, default="f_train.csv", help="path to train dataset")
    parser.add_argument("--test_path", type=str, default="f_test.csv", help="path to test dataset")
    opt = parser.parse_args()

    model = Model(opt.train_path, opt.test_path)
    print('训练集维度:{}\n测试集维度:{}'.format(model.X_train.shape, model.X_test.shape))
    f1_score_train = model.train()
    print('f1_score_train={:.6f}'.format(f1_score_train))
    f_predict = model.predict()
    f_predict.to_csv('f_predict.csv')
